[None][feat] Ring-buffer reuse of out-of-window SWA pages during decode in KV cache manager v2#16288
Draft
lishicheng1996-nv wants to merge 2 commits into
Draft
Conversation
Contributor
📝 WalkthroughWalkthroughAdds an experimental SWA decode slot-reuse configuration, wires it through KV cache manager construction, recycles eligible stale pages during resize, clears parked pages on lifecycle transitions, and tests allocator traffic with the option enabled and disabled. ChangesSWA decode slot reuse
Estimated code review effort: 3 (Moderate) | ~30 minutes Sequence Diagram(s)sequenceDiagram
participant KVCacheManagerV2
participant _KVCache
participant StorageManager
participant SWAReusePocket
KVCacheManagerV2->>_KVCache: resize with reuse enabled
_KVCache->>SWAReusePocket: consume parked page slots
_KVCache->>StorageManager: request remaining GPU slots
_KVCache->>SWAReusePocket: park eligible stale orphaned pages
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py (1)
658-712: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSolid steady-state coverage; consider adding pocket-lifecycle edge cases in
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py.Coverage for the primary decode steady-state (disabled vs. enabled allocator traffic + refcheck correctness) is sufficient and the math checks out. However, the PR's own safety claims aren't directly exercised here:
suspend()/close()draining_swa_reuse_pocket(e.g., suspend mid-decode with a warmed-up pocket, then resume and verify the next crossing allocates fresh rather than reusing a stale/invalid pocket entry).- The 2-entries-per-life-cycle pocket cap (e.g., a resize() that stales two blocks in one call, such as a larger history_length jump).
OutOfPagesErrorrollback interaction with the pocket (PR description states rollback is unchanged, but no test asserts pocket state survives a failedresize()untouched).These are follow-up-worthy rather than blocking, given the core mechanism is already well-covered.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py` around lines 658 - 712, Extend the KV cache manager tests around _run_uncommitted_decode to cover pocket lifecycle safety: suspend and close must drain _swa_reuse_pocket, and after resume the next block crossing must allocate fresh storage rather than reuse stale entries. Add a resize scenario that stales two blocks in one call to verify the two-entry-per-lifecycle cap. Add an OutOfPagesError resize test asserting rollback leaves the pocket state unchanged.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py`:
- Around line 888-905: Revalidate each pocketed holder’s page immediately before
reuse in the `_swa_reuse_pocket` pop path, including `cache_level` and
`has_valid_slot` (and existing commitment checks), before calling
`move_to_new_slot()`. Remove stale or no-longer-eligible holders instead of
reusing them, or otherwise pin them against migration so pocketed pages remain
valid until reuse.
---
Nitpick comments:
In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py`:
- Around line 658-712: Extend the KV cache manager tests around
_run_uncommitted_decode to cover pocket lifecycle safety: suspend and close must
drain _swa_reuse_pocket, and after resume the next block crossing must allocate
fresh storage rather than reuse stale entries. Add a resize scenario that stales
two blocks in one call to verify the two-entry-per-lifecycle cap. Add an
OutOfPagesError resize test asserting rollback leaves the pocket state
unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 7c06be1f-ab1e-49b8-b52b-8a2d1ccc7082
📒 Files selected for processing (6)
tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyitensorrt_llm/runtime/kv_cache_manager_v2/_config.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py
lishicheng1996-nv
marked this pull request as draft
July 13, 2026 03:36
lishicheng1996-nv
changed the title
lishicheng1996-nv
force-pushed
the
feat/swa-decode-ring-reuse
branch
from
July 13, 2026 08:33
75943b8 to
12a194e
Compare
…de in KV cache manager v2 During decode with a sliding window, every tokens_per_block tokens retire one out-of-window block and allocate one new block. Today the staled page takes a full free-list round-trip (release with a ready event, later re-allocation with cross-stream event bookkeeping) while the new block draws a different slot from the allocator. This change adds an opt-in per-sequence reuse pocket (kv_cache_config.enable_swa_ring_reuse, default off): a resize() that orphans a freshly staled page parks it, and a later resize() that needs a new block consumes the parked page's slot directly. On the same CUDA stream the recycled slot is immediately safe to overwrite, so for a 131-token window the sequence's SWA footprint is served by the same ~3 physical blocks in rotation -- a ring buffer over the window's physical pages. Only fully-orphaned pages are parked: uncommitted, with committing disallowed for the cache (so the block radix tree cannot retain them) and GPU-resident. Committed pages kept for prefix reuse are never touched. The pocket is bounded (2 pages per life cycle), consumed and refilled only on resize() success paths so the OutOfPagesError revert is unaffected, and drained on suspend()/close(). Adds TestSwaRingReuse: a disaggregated-decode-style uncommitted generation with FakeEngine refcheck at every step plus deterministic allocator traffic assertions (flag off: one allocation per boundary crossing per life cycle; flag on: the windowed life cycle allocates only until its first page is parked). Signed-off-by: Shicheng Li <shicli@nvidia.com>
lishicheng1996-nv
force-pushed
the
feat/swa-decode-ring-reuse
branch
2 times, most recently
from
July 14, 2026 07:25
04aca57 to
ec4e718
Compare
…t reuse Addresses review: with lower cache levels configured, unlocking a stale page that is still held schedules it for eviction (HELD pages below the last cache level are evictable), so a parked page could be migrated off GPU between park and consume, invalidating the park-time checks. - At park, remove the page from the eviction controller (same idiom as _PageHolder.lock()); migrating a pocketed page would waste bandwidth on dead data anyway. - At consume, re-validate parked pages (uncommitted, valid slot, GPU-resident) and drop any that no longer qualify, keeping the consume path locally correct under future eviction-policy changes. Signed-off-by: Shicheng Li <shicli@nvidia.com>
lishicheng1996-nv
force-pushed
the
feat/swa-decode-ring-reuse
branch
from
July 14, 2026 07:29
ec4e718 to
070fe6b
Compare
lishicheng1996-nv
added
the
api-compatible
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
During decode with a sliding-window-attention (SWA) layer, every
tokens_per_blocktokens the window slides across a block boundary: one out-of-window block is retired and one new block is allocated. Today the staled page takes a full free-list round-trip (release with a ready event, later re-allocation with cross-stream event bookkeeping) while the new block draws a different slot from the global allocator.This PR makes the sliding window's KV storage behave as a ring buffer during decode: at steady state a sequence's SWA footprint is served by the same ~3 physical blocks in rotation (for a 131-token window), with the newest tokens overwriting the block that just left the window — instead of an allocate-new/free-old round-trip per boundary crossing.
It is implemented as an opt-in per-sequence reuse pocket in KV cache manager v2 (
kv_cache_config.enable_swa_ring_reuse, default off): aresize()that orphans a freshly staled page parks it; a laterresize()that needs a new block consumes the parked page's slot directly. On the same CUDA stream the recycled slot is immediately safe to overwrite, so e.g. for a 131-token window a sequence's SWA footprint is served by the same ~3 physical blocks in rotation instead of continuous allocate-then-free.Safety gating (behavior is unchanged except for which physical slot backs the new block):
resize()success paths, so theOutOfPagesErrorrevert path (_lock_held_blocks) is unaffected;suspend()/close().Plumbed through
KvCacheConfig.enable_swa_ring_reuse(prototype) -> pyexecutorKVCacheManagerV2->KVCacheManagerConfig.Test Coverage
New
TestSwaRingReuseintests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py: drives a disaggregated-decode-style uncommitted generation across 8 window advances with FakeEngine refcheck at every step, plus deterministic allocator-traffic assertions:2 x crossings);crossings + 1).Existing
TestNoBatchingbattery passes unchanged under the patched module (flag off is a no-op: two cheap branches).Perf A/B on DeepSeek-V4 disaggregated GEN (GB200, dep32/EP32, 256K ISL, per-rank batch 84, MTP3, MoE backend MEGAMOE_DEEPGEMM, 2688/2688 requests completed on both arms):
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.Summary by CodeRabbit
New Features
Bug Fixes
Tests